import requests
import json
import time
import os
import sys
import urllib3
import json
import argparse
urllib3.disable_warnings()
# os.environ['http_proxy'] = '127.0.0.1:8888'
# os.environ['https_proxy'] = '127.0.0.1:8888'
cert_path = "C:\\Users\\72189894\\adproxy.crt"

class ChatBot:
    def __init__(self, api_url, token):
        self.api_url = api_url
        self.token = token
        self.conversation_id = ''  # 使用时间戳生成初始会话ID
    
    def send_query(self, query):
        """发送查询并处理流式响应"""
        print(f"Sending query: {query}")
        payload = {
            "query": query,
            "requestExtendInfo": "{\"serverType\":0,\"isOnline\":0,\"modelName\":\"blue-deekseek-r1\"}",
            "modelKey": "blue-deekseek-r1",
            "source": "",
            "machineId": "",
            "ideVersion": "",
            "pluginVersion": "",
            "userId": "72189894",
            "sessionId": "",
            "questionDisplay": "",
            "currentOpenFilePath": "",
            "currentOpenFileContent": "",
            "conversationId": self.conversation_id
        }
        
        headers = {
            'Content-Type': 'application/json',
            'token': self.token  # 添加token到请求头
        }
        
        try:
            response = requests.post(
                self.api_url,
                data= json.dumps(payload),
                headers=headers,
                proxies = {
                    "http": "http://skyproxy.vivo.xyz:80",
                    "https": "http://skyproxy.vivo.xyz:80"
                },
                verify=cert_path,
                stream=True,
                timeout=300
            )
            
            if response.status_code != 200:
                print(f"请求失败，状态码: {response.status_code}, 响应: {response.text}")
                return None
            
            full_response = ""
            for line in response.iter_lines():
                if line:
                    decoded_line = line.decode('utf-8')
                    if decoded_line.startswith('data:'):
                        data_str = decoded_line[5:].strip()
                        
                        if data_str == '[DONE]':
                            break
                        
                        try:
                            data = json.loads(data_str)
                            answer_chunk = data.get('answer', '')
                            self.conversation_id = data.get('conversation_id', '')
                            print(answer_chunk, end='', flush=True)  # 实时打印流式响应
                            full_response += answer_chunk
                        except json.JSONDecodeError:
                            print(f"\nJSON解析错误: {data_str}")
            
            print()  # 打印换行符使输出整洁
            return full_response
        
        except requests.exceptions.RequestException as e:
            print(f"请求异常: {e}")
            return None

if __name__ == "__main__":
    # 替换为你的实际API URL和token
    API_URL = "https://ai-code-chat.vmic.xyz/api/generate/v2/chat/stream"
    TOKEN = "niM6bpUjc.s6xEgwb9lyXChL4lipNbR1o8rQjuZLceajmbEXazGwe_xgLdCAsV.5HvQdCKJzSTs*"
    parser = argparse.ArgumentParser(description="Chat with the AI")
    parser.add_argument('-t', '--token', help='Your token', required=False)
    args = parser.parse_args()
    if args.token != None:
        TOKEN = args.token
        # write token to file
        with open('.blue_token.txt', 'w') as f:
            f.write(TOKEN)
    elif os.path.exists('.blue_token.txt'):
        try:
            with open('.blue_token.txt', 'r') as f:
                TOKEN = f.read().strip()
                print("Token loaded from .blue_token.txt")
        except Exception as e:
            print(f"Error loading token from .blue_token.txt: {e}")
    else:
        print("No token provided or found.")
    print("Token:", TOKEN)
    bot = ChatBot(API_URL, TOKEN)
    print("聊天机器人已启动 (输入 'exit' 退出)")
    while True:
        user_input = input("\n你: ")
        if user_input.lower() == 'exit':
            break
        
        if user_input.strip():
            print("AI: ", end='', flush=True)
            bot.send_query(user_input)